package aes

import (
	"crypto/aes"
	"crypto/cipher"
	"fmt"
	"github.com/pkg/errors"
)

type CryptoECB struct {
	block       cipher.Block
	key         []byte
	size        int
	paddingMode PaddingMode
	way         Way
}

func newAESCryptoECB(key []byte, mode PaddingMode, way Way) (Crypto, error) {
	b, err := aes.NewCipher(key)

	if err != nil {
		return nil, err
	}

	r := &CryptoECB{
		block:       b,
		key:         key,
		size:        b.BlockSize(),
		paddingMode: Pkcs7Padding,
		way:         way,
	}

	if mode != "" {
		r.paddingMode = mode
	}

	return r, nil
}

func (a *CryptoECB) Encrypt(plainText []byte) (cipherText []byte, err error) {
	defer func() {
		if e := recover(); e != nil {
			err = errors.Wrap(errors.New("aes ecb encrypt error"), fmt.Sprintf("%+v", e))
			return
		}
	}()

	switch a.paddingMode {
	case ZeroPadding:
		plainText = zeroPadding(plainText, a.size)
	case Pkcs5Padding:
		plainText = pkcs5Padding(plainText)
	case Pkcs7Padding:
		plainText = pkcs7Padding(plainText, a.size)
	}

	// 将明文字符串按照初始块的长度均分为若干段，分别进行加密组成加密切片返回
	cipherText = make([]byte, len(plainText))
	for bs, be := 0, a.size; bs < len(plainText); bs, be = bs+a.size, be+a.size {
		a.block.Encrypt(cipherText[bs:be], plainText[bs:be])
	}

	return Encode(cipherText, a.way), nil
}

func (a *CryptoECB) Decrypt(cipherText []byte) (plainText []byte, err error) {
	defer func() {
		if e := recover(); e != nil {
			err = errors.Wrap(errors.New("aes ecb decrypt error"), fmt.Sprintf("%+v", e))
			return
		}
	}()

	cipherText, err = Decode(cipherText, a.way)
	if err != nil {
		return nil, errors.WithStack(err)
	}

	plainText = make([]byte, len(cipherText))
	for bs, be := 0, a.size; bs < len(cipherText); bs, be = bs+a.size, be+a.size {
		a.block.Decrypt(plainText[bs:be], cipherText[bs:be])
	}

	if a.paddingMode == ZeroPadding {
		plainText = zeroUnPadding(plainText)
	} else {
		plainText = pkcsUnPadding(plainText)
	}

	return plainText, nil
}
